import torch
import torchvision
from torch.utils.data import DataLoader
from module import resnet18

# 数据集路径
root = "../dataset"
# 批加载图片数量
batch_size = 20

# 测试集
test_set = torchvision.datasets.CIFAR100(root=root, train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
# loader
test_dataloader = DataLoader(test_set, batch_size=batch_size)
# 数据集长度
test_set_len = len(test_set)

# 创建网络模型
resnet = resnet18()
# 载入权重
resnet.load_state_dict(torch.load("../Module/resnet18.pth"))
print("加载网络模型成功")

print("开始测试")
resnet.eval()
total_accuracy = 0  # 总的测试正确数
test_step = 0  # 记录测试次数

with torch.no_grad():
    for data in test_dataloader:
        imgs, targets = data
        output = resnet(imgs)

        accuracy = (output.argmax(1) == targets).sum()
        total_accuracy += accuracy

        test_step += 1
        if test_step % 100 == 0:
            print("测试次数：{} \taccuracy：{}/{}".format(test_step, total_accuracy, test_set_len))
print("测试集上的accuracy：{}%".format(100 * total_accuracy / test_set_len))
